import numpy as np
import glob
from PIL import Image
import matplotlib.pyplot as plt


def file_arr2png(in_file_name, out_file_name, dtype, crop_out_pml, dim_x, dim_y, pml_x, pml_y):
    data = np.fromfile(in_file_name, dtype=dtype)
    data = np.reshape(data, newshape=(dim_y, dim_x))
    if crop_out_pml:
        data = data[pml_y[0]:(dim_y - pml_y[1]), pml_x[0]:(dim_x - pml_x[1])]

    im = Image.fromarray(data)
    im = im.convert('L')
    im = im.rotate(90, expand=1)
    im = im.transpose(Image.FLIP_TOP_BOTTOM)
    im.save(out_file_name)


def slice3d2png(in_file_name, out_file_name, dtype, dim_x, dim_y, dim_z, target_dim, target_value):
    data = np.fromfile(in_file_name, dtype=dtype)
    data = np.reshape(data, newshape=(dim_z, dim_y, dim_x))

    if target_dim == 'x':
        data_out = data[:, :, target_value]
    elif target_dim == 'y':
        data_out = data[:, target_value, :]
    elif target_dim == 'z':
        data_out = data[target_value, :, :]
    else:
        print("Unrecognized value of 'target_dim': ", target_dim)
        exit(1)

    im = Image.fromarray(data_out)
    im = im.convert('L')
    if target_dim == 'z':
        im = im.rotate(90, expand=1)
        im = im.transpose(Image.FLIP_TOP_BOTTOM)
    im.save(out_file_name)


def draw_wave_curve(file):                  # to plot the wave form of certain output points
    data = np.fromfile(file, np.float64)
    plt.plot(data)
    # plt.show()                            # one can disable the annotation if one wants to view this figure on the screen
    plt.savefig(str.replace(file, ".dat", ".png"))


if __name__ == '__main__':

    file_list = glob.glob("/home/lactose/Sea_pml_test/lactose_PWave_2021.4.18/Sea_small_server/*_x*_t*.dat")

    for file in file_list:
        file_arr2png(in_file_name=file, out_file_name=str.replace(file, ".dat", ".png"), dtype=np.int16, crop_out_pml=True,
                     dim_x=1051, dim_y=601, pml_x=(50, 0), pml_y=(50, 50))

    draw_wave_curve("/home/lactose/Sea_pml_test/lactose_PWave_2021.4.18/Sea_small_server/Sea_small_server_out_pt_id0_x25_y250_z250.dat")
